import sys, os, copy, math
from topology_engineer import *
sys.path.append("..")
#import aurora_network
from gurobipy import *

'''		
Design fractional topology based on 
traffic clusters using LP formulation
'''
#combines multiple fractional topology into one
class IdealTopologyEngineer(TopologyEngineer):
	def __init__(self, aurora_network, all_paths, all_traffic_snapshots):
		num_snapshots = len(all_traffic_snapshots)
		TopologyEngineer.__init__(self, aurora_network, num_snapshots)
		self.toe_class = TopologyEngineeringClass.IDEAL
		## NOTE: have to rezero the cached logical topologies here because when we initialize the parent class,
		## 		it will assume that the topology is static (see TopologyEngineer class above).
		self.all_traffic_snapshots = all_traffic_snapshots
		self.cached_logical_topologies = {}
		self.all_paths = all_paths
		return

	def get_logical_topology_at_time(self, time_index):
		adj_matrix = None
		assert(time_index < self.num_snapshots)
		adj_matrix = self._topology_engineer(self.all_traffic_snapshots[time_index], self.all_paths)
		return self.aurora_network.round_fractional_topology_giant_switch(adj_matrix, [])

	def get_filename_param(self):
		return "perfect"

	# returns a string description of this class
	def get_string(self):
		return "Perfect ToE"

	def is_static(self):
		return False

	# given a scale up factor, figure out the routing that minimizes two-hop paths if they exists
	## Todo(jason) : fix this and debug
	def _minimize_multihop(self, traffic_matrix, scaling_factor, all_paths):
		## using QP to reduce multihop reliance
		nblocks = self.aurora_network.get_num_blocks()
		model = Model("minimize multihop")
		model.setParam( 'OutputFlag', False )
		interpod_link_counts = [None] * nblocks
		routing_weights = {}
		link_capacity_constraints = [None] * nblocks
		## setup the link constraints and all optization variables, 
		## which are the link counts and routing variables
		objective_function = QuadExpr()
		for i in range(nblocks):
			interpod_link_counts[i] = [None] * nblocks
			link_capacity_constraints[i] = [None] * nblocks
			for j in range(nblocks):
				if (i != j):
					link_capacity_constraints[i][j] = LinExpr()
					upper_bound = min(self.aurora_network.get_num_links(i), self.aurora_network.get_num_links(j))
					interpod_link_counts[i][j] = model.addVar(lb=0, ub=upper_bound, obj=0, vtype=GRB.CONTINUOUS, name="lc" + str(i) + ":" + str(j))
					for path in all_paths[i][j]:
						var = model.addVar(lb=0, ub=GRB.INFINITY, obj=0, vtype=GRB.CONTINUOUS, name="w_{}".format(path))
						routing_weights[path] = var
						# if multihop path, then add to objective function
						if len(path) > 2: 
							objective_function.add(var * var, mult=1.)

		## add radix degree constraints
		for pod in range(nblocks):
			row_constraint = LinExpr()
			col_constraint = LinExpr()
			nlinks = self.aurora_network.get_num_links(pod)
			for target_pod in range(nblocks):
				if target_pod != pod:
					row_constraint.add(interpod_link_counts[pod][target_pod], mult=1.)
					col_constraint.add(interpod_link_counts[target_pod][pod], mult=1.)
			model.addConstr(lhs=row_constraint, sense=GRB.LESS_EQUAL, rhs=nlinks)
			model.addConstr(lhs=col_constraint, sense=GRB.LESS_EQUAL, rhs=nlinks)

		## add achievability of flows for all i j pairs
		for i in range(nblocks):
			for j in range(nblocks):
				if (i != j):
					flow_achievability_constr = LinExpr()
					for path in all_paths[i][j]:
						flow_achievability_constr.add(routing_weights[path], 1.)
						path_len = len(path)
						curr_node = path[0]
						for next_node_index in range(1, path_len, 1):
							next_node = path[next_node_index]
							link_capacity_constraints[curr_node][next_node].add(routing_weights[path], mult=1.)
							curr_node = next_node
					model.addConstr(lhs=scaling_factor * traffic_matrix[i][j], sense=GRB.LESS_EQUAL, rhs=flow_achievability_constr)
		
		## add link capacity limit for all paths constraints
		for i in range(nblocks):
			for j in range(nblocks):
				if (i != j):
					capacity = min(self.aurora_network.get_link_capacity(i), self.aurora_network.get_link_capacity(i))
					model.addConstr(lhs=link_capacity_constraints[i][j], sense=GRB.LESS_EQUAL, rhs=interpod_link_counts[i][j] * capacity)
		# set up the objective function
		model.setObjective(objective_function, GRB.MINIMIZE)
		# start optimizing
		adj_matrix = np.zeros( (nblocks, nblocks,) )
		try: 
			model.optimize()
			for i in range(nblocks):
				for j in range(nblocks):
					if (i != j):
						adj_matrix[i][j] = interpod_link_counts[i][j].x
			#print adj_matrix
			return adj_matrix
		except GurobiError as e:
			print ("MinimizeMultihop: Error code " + str(e. errno ) + ": " + str(e))
			return None
		except AttributeError as e1:
			print ("MinimizeMultihop: Encountered an attribute error " + ": " + str(e1))
			return None
	
	# first step of topology design which solves the LP that maximizes traffic scale up
	def _maximize_traffic_scaleup(self, traffic_matrix, all_inter_block_paths):
		nblocks = self.aurora_network.get_num_blocks()
		model = Model("maximize traffic scale up")
		model.setParam( 'OutputFlag', False )
		mu = model.addVar(lb=0, ub=GRB.INFINITY, obj=1., vtype=GRB.CONTINUOUS, name="mu")
		fractional_topology_var = [None] * nblocks
		routing_weights = {}
		link_capacity_constraints = [0] * nblocks
		## setup the link constraints and all optization variables, 
		## which are the link counts and routing variables
		for i in range(nblocks):
			fractional_topology_var[i] = [None] * nblocks
			link_capacity_constraints[i] = [None] * nblocks
			for j in range(nblocks):
				if i != j:
					link_capacity_constraints[i][j] = LinExpr()
					upper_bound = min(self.aurora_network.get_num_links(i), self.aurora_network.get_num_links(j))
					fractional_topology_var[i][j] = model.addVar(lb=0, ub=upper_bound, obj=0, vtype=GRB.CONTINUOUS, name="lc" + str(i) + ":" + str(j))
					for path in all_inter_block_paths[i][j]:
						var = model.addVar(lb=0, ub=GRB.INFINITY, obj=0, vtype=GRB.CONTINUOUS, name="w_{}".format(path))
						routing_weights[path] = var

		## add radix degree constraints
		for pod in range(nblocks):
			row_constraint = LinExpr()
			col_constraint = LinExpr()
			for target_pod in range(nblocks):
				if target_pod != pod:
					row_constraint.add(fractional_topology_var[pod][target_pod], mult=1.)
					col_constraint.add(fractional_topology_var[target_pod][pod], mult=1.)
			model.addConstr(lhs=row_constraint, sense=GRB.LESS_EQUAL, rhs=float(self.aurora_network.get_num_links(pod)))
			model.addConstr(lhs=col_constraint, sense=GRB.LESS_EQUAL, rhs=float(self.aurora_network.get_num_links(pod)))

		## add achievability of flows for all i j pairs
		for i in range(nblocks):
			for j in range(nblocks):
				if (i != j):
					flow_achievability_constr = LinExpr()
					for path in all_inter_block_paths[i][j]:
						flow_achievability_constr.add(routing_weights[path], 1.)
						path_len = len(path)
						curr_node = path[0]
						for next_node_index in range(1, path_len, 1):
							next_node = path[next_node_index]
							link_capacity_constraints[curr_node][next_node].add(routing_weights[path], 1.)
							curr_node = next_node
					model.addConstr(lhs=mu * traffic_matrix[i][j], sense=GRB.LESS_EQUAL, rhs=flow_achievability_constr)
		
		## add link capacity limit for all paths constraints
		for i in range(nblocks):
			for j in range(nblocks):
				if (i != j):
					capacity = min(self.aurora_network.get_link_capacity(i),self.aurora_network.get_link_capacity(j))
					model.addConstr(lhs=link_capacity_constraints[i][j], sense=GRB.LESS_EQUAL, rhs=fractional_topology_var[i][j] * capacity)

		# set up the objective function
		model.setObjective(mu, GRB.MAXIMIZE)
		# start optimizing
		try: 
			model.optimize()
			adj_matrix = np.zeros((nblocks, nblocks))
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						adj_matrix[i][j] = fractional_topology_var[i][j].x
			return mu.x, adj_matrix # returns the max scale up factor
		except GurobiError as e:
			print ("MaximizeTrafficScaleUp: Error code " + str(e. errno ) + ": " + str(e))
			return None
		except AttributeError :
			print ("MaximizeTrafficScaleUp: Encountered an attribute error ")
			return None

	# conducts topology engineering
	# aurora network - Aurora Network, contains all the parameters
	# training_traffic_snapshots - all the snapshots used for training the topology
	# number of k - which is the number of traffic clusters we use to solve engineer topology
	def _topology_engineer(self, traffic_snapshot, all_paths):
		traffic_matrix = self.aurora_network.get_traffic_matrix_from_traffic_snapshot(traffic_snapshot)
		mu, _ = self._maximize_traffic_scaleup(traffic_matrix, all_paths)
		adj_matrix_stage_2 = None
		relaxation = 1.
		while adj_matrix_stage_2 is None and relaxation > 0:
			adj_matrix_stage_2 = self._minimize_multihop(traffic_matrix, relaxation * mu, all_paths)
			relaxation -= 0.01
		adj_matrix = self._reinforce_fractional_topology(adj_matrix_stage_2)
		if adj_matrix is None:
			adj_matrix = adj_matrix_stage_2
		return adj_matrix

	# conducts topology engineering
	# aurora network - Aurora Network, contains all the parameters
	# training_traffic_snapshots - all the snapshots used for training the topology
	# number of k - which is the number of traffic clusters we use to solve engineer topology
	def topology_engineer_given_TM(self, traffic_matrix, all_paths):
		mu, _ = self._maximize_traffic_scaleup(traffic_matrix, all_paths)
		adj_matrix_stage_2 = None
		relaxation = 1.
		while adj_matrix_stage_2 is None and relaxation > 0:
			adj_matrix_stage_2 = self._minimize_multihop(traffic_matrix, relaxation * mu, all_paths)
			relaxation -= 0.01
		adj_matrix = self.progressive_filling(adj_matrix_stage_2)
		if adj_matrix is None:
			adj_matrix = adj_matrix_stage_2
		return adj_matrix